/**
 * Trusted application
 */
#include <fcntl.h>
#include <stdlib.h>
#include <tee_ext_api.h>
#include <tee_log.h>
#include <securec.h>
#include <unistd.h>
#include <sys/stat.h>
#include <openssl/evp.h>
#include <openssl/err.h>
#include "infer.h"

#define PARAM_COUNT      4
#define IN_BUFFER_INDEX 2
#define OUT_BUFFER_INDEX 3
#define BLOCK_SIZE 8192
#define MODEL_FILE_NAME_MAXLEN 32
#define DECRYPTED_FILE_NAME "./llm_model.gguf"
#define MODEL_FILE_MAXSIZE 512 * 1024 * 1024

enum {
    CMD_GET_TA_VERSION = 1,
    CMD_BEGIN_SEND_MODEL_FILE = 2,
    CMD_SEND_MODEL_FILE_SEGMENT = 3,
    CMD_INFERENCE = 4,
    CMD_DELETE_MODEL_FILE = 5,
};
char model_file_name[64] = {0};
int model_file_fd = 0;
int model_file_size = 0;
int current_receive_file_size = 0;
unsigned char my_key[256] = {0xA, 0xB, 0xC, 0xD, 0xE, 0xF, 0x0, 0x1};
unsigned char initial_vector[128] = {0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0x10};
EVP_CIPHER_CTX* ctx = NULL;

/* MD5 Calculation */
#ifdef MD5_CHECK


#include <string.h>
#include <stdint.h>

#define F(b, c, d) (((b) & (c)) | ((~b) & (d)))
#define G(b, c, d) (((b) & (d)) | ((c) & (~d)))
#define H(b, c, d) ((b) ^ (c) ^ (d))
#define I(b, c, d) ((c) ^ ((b) | (~d)))

#define LEFTROTATE(num, n) (((num) << n) | ((num >> (32 - n))))

const uint32_t T[64] = { 0xd76aa478, 0xe8c7b756, 0x242070db, 0xc1bdceee,
						0xf57c0faf, 0x4787c62a, 0xa8304613, 0xfd469501,
						0x698098d8, 0x8b44f7af, 0xffff5bb1, 0x895cd7be,
						0x6b901122, 0xfd987193, 0xa679438e, 0x49b40821,
						0xf61e2562, 0xc040b340, 0x265e5a51, 0xe9b6c7aa,
						0xd62f105d, 0x02441453, 0xd8a1e681, 0xe7d3fbc8,
						0x21e1cde6, 0xc33707d6, 0xf4d50d87, 0x455a14ed,
						0xa9e3e905, 0xfcefa3f8, 0x676f02d9, 0x8d2a4c8a,
						0xfffa3942, 0x8771f681, 0x6d9d6122, 0xfde5380c,
						0xa4beea44, 0x4bdecfa9, 0xf6bb4b60, 0xbebfbc70,
						0x289b7ec6, 0xeaa127fa, 0xd4ef3085, 0x04881d05,
						0xd9d4d039, 0xe6db99e5, 0x1fa27cf8, 0xc4ac5665,
						0xf4292244, 0x432aff97, 0xab9423a7, 0xfc93a039,
						0x655b59c3, 0x8f0ccc92, 0xffeff47d, 0x85845dd1,
						0x6fa87e4f, 0xfe2ce6e0, 0xa3014314, 0x4e0811a1,
						0xf7537e82, 0xbd3af235, 0x2ad7d2bb, 0xeb86d391 };

const uint32_t S[64] = { 7, 12, 17, 22, 7, 12, 17, 22, 7, 12, 17, 22, 7, 12, 17, 22,
						 5,  9, 14, 20, 5,  9, 14, 20, 5,  9, 14, 20, 5,  9, 14, 20,
						 4, 11, 16, 23, 4, 11, 16, 23, 4, 11, 16, 23, 4, 11, 16, 23,
						 6, 10, 15, 21, 6, 10, 15, 21, 6, 10, 15, 21, 6, 10, 15, 21 };

void int2byte(uint32_t val, uint8_t *bytes)
{
	bytes[0] = (uint8_t)val;
	bytes[1] = (uint8_t)(val >> 8);
	bytes[2] = (uint8_t)(val >> 16);
	bytes[3] = (uint8_t)(val >> 24);
}

uint32_t byte2int(const uint8_t *bytes)
{
	return (uint32_t)bytes[0]
		| ((uint32_t)bytes[1] << 8)
		| ((uint32_t)bytes[2] << 16)
		| ((uint32_t)bytes[3] << 24);
}

int MD5(const char* filepath, uint8_t *result) {
	FILE *fp = NULL;
	uint8_t buffer[64];
	uint8_t* temp = NULL;
	size_t count = 0, offset, i;
	uint32_t X[16];
	int flag = 0;

	if ((fp = fopen(filepath, "rb+")) == NULL) {
		tloge("[ERROR] File in %s not found.", filepath);
		return 0;
	}

	uint32_t A, B, C, D;
	A = 0x67452301;
	B = 0xEFCDAB89;
	C = 0x98BADCFE;
	D = 0X10325476;

	while (!feof(fp)) {
		memset(buffer, 0, sizeof(buffer));
		size_t len = fread(buffer, 1, 64, fp);
		count += len;
		if (feof(fp)) {
			flag = 1;

			size_t new_len;
			for (new_len = len + 1; new_len % 64 != 56; new_len++)
				;

			temp = (uint8_t*)malloc(new_len + 8);
			memcpy(temp, buffer, len);

			temp[len] = 0x80;
			for (offset = len + 1; offset < new_len; offset++)
				temp[offset] = 0;

			int2byte(count * 8, temp + new_len);
			int2byte(count >> 29, temp + new_len + 4);
			len = new_len;
		}

		for (offset = 0; offset < len; offset += 64) {
			if (flag == 1) {
				memcpy(buffer, temp + offset, 64);
			}

			for (int i = 0; i < 16; i++) {
				X[i] = byte2int(buffer + i * 4);
			}

			uint32_t a, b, c, d, temp, g, k;
			a = A;
			b = B;
			c = C;
			d = D;

			for (i = 0; i < 64; i++) {
				if (i < 16) {
					g = F(b, c, d);
					k = i;
				}
				else if (i < 32) {
					g = G(b, c, d);
					k = (1 + 5 * i) % 16;
				}
				else if (i < 48) {
					g = H(b, c, d);
					k = (5 + 3 * i) % 16;
				}
				else {
					g = I(b, c, d);
					k = (7 * i) % 16;
				}
				temp = d;
				d = c;
				c = b;
				b = b + LEFTROTATE((a + g + X[k] + T[i]), S[i]);
				a = temp;
			}

			A += a;
			B += b;
			C += c;
			D += d;

		}
	}
	
	free(temp);
    fclose(fp);

	int2byte(A, result);
	int2byte(B, result + 4);
	int2byte(C, result + 8);
	int2byte(D, result + 12);
    tloge("%s MD5 result:%2.2x%2.2x%2.2x%2.2x%2.2x%2.2x%2.2x%2.2x%2.2x%2.2x"
            "%2.2x%2.2x%2.2x%2.2x%2.2x%2.2x\n", filepath, result[0], result[1],
            result[2], result[3], result[4], result[5], result[6], result[7],
            result[8], result[9], result[10], result[11], result[12],
            result[13], result[14], result[15]);
	return 1;
}

#endif // MD5_CHECK

/**
 * Function TA_CreateEntryPoint
 * Description:
 *   The function TA_CreateEntryPoint is the Trusted Application's constructor,
 *   which the Framework calls when it creates a new instance of this Trusted Application.
 */
TEE_Result TA_CreateEntryPoint(void)
{
    TEE_Result ret;

    tlogd("----- TA entry point ----- ");

    ret = AddCaller_CA_exec("/vendor/bin/teec_hello", 0);
    if (ret == TEE_SUCCESS) {
        tlogd("TA entry point: add ca whitelist success");
    } else {
        tloge("TA entry point: add ca whitelist failed");
        return TEE_ERROR_GENERIC;
    }

    return TEE_SUCCESS;
}

/**
 * Function TA_OpenSessionEntryPoint
 * Description:
 *   The Framework calls the function TA_OpenSessionEntryPoint
 *   when a client requests to open a session with the Trusted Application.
 *   The open session request may result in a new Trusted Application instance
 *   being created.
 */
TEE_Result TA_OpenSessionEntryPoint(uint32_t parm_type,
    TEE_Param params[PARAM_COUNT], void** session_context)
{
    (void)parm_type;
    (void)params;
    (void)session_context;
    tlogd("---- TA open session -------- ");

    return TEE_SUCCESS;
}

/**
 * Function TA_InvokeCommandEntryPoint
 * Description:
 *   The Framework calls this function when the client invokes a command
 *   within the given session.
 */
TEE_Result TA_InvokeCommandEntryPoint(void* session_context, uint32_t cmd,
    uint32_t parm_type, TEE_Param params[PARAM_COUNT])
{
    TEE_Result ret;
    (void)session_context;
    unsigned long free_mem = 0;
    int min_buffer;

#ifdef MD5_CHECK
    uint8_t hash_result[16];
#endif


    tlogd("---- TA invoke command ----------- ");
    switch (cmd) {
    case CMD_BEGIN_SEND_MODEL_FILE:
        if (!check_param_type(parm_type,
            TEE_PARAM_TYPE_VALUE_INPUT,// model file size in byte
            TEE_PARAM_TYPE_MEMREF_INPUT,// model file name
            TEE_PARAM_TYPE_NONE, TEE_PARAM_TYPE_NONE)) {
            tloge("Bad expected parameter types");
            return TEE_ERROR_BAD_PARAMETERS;
        }
        tloge("ready to call openssl\n");
        OPENSSL_init_crypto(OPENSSL_INIT_ADD_ALL_CIPHERS \
         | OPENSSL_INIT_ADD_ALL_DIGESTS, NULL);
        OPENSSL_init_crypto(OPENSSL_INIT_LOAD_CRYPTO_STRINGS, NULL);
        ctx = EVP_CIPHER_CTX_new();
        EVP_CIPHER_CTX_init(ctx);
        if (!EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, my_key,
            initial_vector)) {
            tloge("TA: EVP_DecryptInit_ex error, not exit for debug\n");
        }
 
        EVP_CIPHER_CTX_set_padding(ctx,EVP_PADDING_PKCS7);

#ifdef MD5_CHECK
        if (MD5("/libllama.so", hash_result) == 1) {
        } else {
            return TEE_ERROR_STORAGE_NOT_AVAILABLE;
        }
        if (MD5("/libggml.so", hash_result) == 1) {
        } else {
            return TEE_ERROR_STORAGE_NOT_AVAILABLE;
        }
        if (MD5("/libgcc_s.so.1", hash_result) == 1) {
        } else {
            return TEE_ERROR_STORAGE_NOT_AVAILABLE;
        }
        if (MD5("/libstdc++.so.6", hash_result) == 1) {
        } else {
            return TEE_ERROR_STORAGE_NOT_AVAILABLE;
        }
        if (MD5("/libinfer.so", hash_result) == 1) {
        } else {
            return TEE_ERROR_STORAGE_NOT_AVAILABLE;
        }
#endif
        if (params[0].value.a > 0 && params[0].value.a <= MODEL_FILE_MAXSIZE) {
            model_file_size = params[0].value.a;
        } else {
            return TEE_ERROR_OUT_OF_MEMORY;
        }
        min_buffer = (int)params[1].memref.size;
        if (min_buffer >= MODEL_FILE_NAME_MAXLEN || min_buffer < 0) {
            min_buffer = MODEL_FILE_NAME_MAXLEN;
        }
        memcpy(model_file_name, params[1].memref.buffer, min_buffer);
        if (model_file_fd > 0) {
            close(model_file_fd);
            tlogi("TA: close model_file_fd before sending model.\n");
            model_file_fd = 0;
        }
        model_file_fd = open(DECRYPTED_FILE_NAME, O_RDWR|O_CREAT|O_TRUNC,
            S_IRUSR|S_IWUSR);
        tloge("TA: open %s returns %d, errno=%d\n", DECRYPTED_FILE_NAME, 
            model_file_fd, errno);
        if (model_file_fd < 0) {
            ret = TEE_ERROR_NOT_SUPPORTED;
        } else {
            ret = TEE_SUCCESS;
        }
        break;
    case CMD_SEND_MODEL_FILE_SEGMENT:
        if (!check_param_type(parm_type,
            TEE_PARAM_TYPE_VALUE_INPUT, // sequence
            TEE_PARAM_TYPE_MEMREF_INPUT, // buffer
            TEE_PARAM_TYPE_NONE,TEE_PARAM_TYPE_NONE)) {
            tloge("Bad expected parameter types");
            return TEE_ERROR_BAD_PARAMETERS;
        }

        min_buffer = (int)params[1].memref.size;
        if (min_buffer >= BLOCK_SIZE || min_buffer < 0) {
            min_buffer = BLOCK_SIZE;
        }
        current_receive_file_size += min_buffer;

        // AES-256-CBC decryption
        // input length + cipher_block_size for EVP_DecryptUpdate
        uint8_t write_buffer[BLOCK_SIZE+32] = {0x00};    
        int toEnBufLen = 0;
        int curEnLen = 0;
        if(!EVP_DecryptUpdate(ctx, write_buffer, &toEnBufLen, 
            (const unsigned char *)(params[1].memref.buffer), min_buffer)){
            tloge("EVP_DecryptUpdate failed!\n");
            return TEE_ERROR_CORRUPT_OBJECT;
        }
        if (min_buffer < BLOCK_SIZE){
            // last block, finish decryption
            if(!EVP_DecryptFinal_ex(ctx, write_buffer + toEnBufLen, &curEnLen)) {
                tloge("EVP_DecryptFinal failed!  \n");
                return TEE_ERROR_CORRUPT_OBJECT;
            }
            ctx = NULL;
            toEnBufLen += curEnLen;
        }

        // write decrypted data to file
        ssize_t write_ret;
        write_ret = write(model_file_fd, write_buffer, toEnBufLen);
        if (write_ret <= 0) {
            tloge("TA: write file returns %llu, errno=%d\n", write_ret, errno);
        }
        if (write_ret < 0) {
            return TEE_ERROR_STORAGE_NO_SPACE;
        }
        ret = TEE_SUCCESS;
        break;
    case CMD_INFERENCE:
#ifdef DEBUG_GET_CHCORE_FREEMEM
        asm volatile(
            "mov x8, %1\n"
            "svc #0\n"
            "mov %0, x0\n"
            : "=r" (free_mem)
            : "i" (222)
            : "x0", "x8"
        );
        tloge("get free mem is %lu\n", free_mem);
#endif
        if (model_file_size != current_receive_file_size) {
            tloge("TA error: model_file_size = %d, current_receive_file_size ="
                " %d\n", model_file_size, current_receive_file_size);
            return TEE_ERROR_CORRUPT_OBJECT;
        }
        sync();

#ifdef MD5_CHECK
        if (MD5(DECRYPTED_FILE_NAME, hash_result) == 1) {
        } else {
            return TEE_ERROR_STORAGE_NOT_AVAILABLE;
        }
#endif

        if (!check_param_type(parm_type,
            TEE_PARAM_TYPE_MEMREF_INPUT, // input buffer
            TEE_PARAM_TYPE_MEMREF_OUTPUT, // output buffer
            TEE_PARAM_TYPE_NONE,TEE_PARAM_TYPE_NONE)) {
            tloge("Bad expected parameter types");
            ret = TEE_ERROR_BAD_PARAMETERS;
            break;
        }

        char *input_buffer = (char*)params[0].memref.buffer;
        size_t input_len = params[0].memref.size;
        char *output_buffer = (char*)params[1].memref.buffer;
        size_t output_len = params[1].memref.size;
        char *argv[] = {"./llama-cli",  "-cnv", "-m", DECRYPTED_FILE_NAME,
            "-p", "You are a clever student.", "-n", "256"};
        infer(8, argv, input_buffer, input_len, output_buffer, output_len);

        ret = TEE_SUCCESS;
        break;
    case CMD_DELETE_MODEL_FILE:
        ret = TEE_SUCCESS;
        break;
    default:
        tloge("Unknown cmd is %u", cmd);
        ret = TEE_ERROR_BAD_PARAMETERS;
    }

    return ret;
}

/**
 * Function TA_CloseSessionEntryPoint
 * Description:
 *   The Framework calls this function to close a client session.
 *   During the call to this function the implementation can use
 *   any session functions.
 */
void TA_CloseSessionEntryPoint(void* session_context)
{
    (void)session_context;
    if (model_file_fd > 0) {
        close(model_file_fd);
    }
    model_file_fd = 0;
    model_file_size = 0;
    current_receive_file_size = 0;
    tlogd("---- close session ----- ");
}

/**
 * Function TA_DestroyEntryPoint
 * Description:
 *   The function TA_DestroyEntryPoint is the Trusted Application's destructor,
 *   which the Framework calls when the instance is being destroyed.
 */
void TA_DestroyEntryPoint(void)
{
    tlogd("---- destroy TA ---- ");
}
